/* Copyright (c) 2023 Renmin University of China
RMDB is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
        http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details. */

#include <algorithm>
#include <functional>

#include "lock_manager.h"
#include "common/config.h"
#include "transaction/transaction_manager.h"

#define ANYLOCK \
    { LockMode::SHARED, LockMode::EXCLUSIVE, LockMode::INTENTION_SHARED, LockMode::INTENTION_EXCLUSIVE, LockMode::S_IX }

std::chrono::milliseconds cycle_detection_interval = std::chrono::milliseconds(50);

/**
 * @description: 申请行级共享锁
 * @return {bool} 加锁是否成功
 * @param {Transaction*} txn 要申请锁的事务对象指针
 * @param {Rid&} rid 加锁的目标记录ID 记录所在的表的fd
 * @param {int} tab_fd
 */
bool LockManager::lock_shared_on_record(Transaction* txn, const Rid& rid, int tab_fd) {
    return lock_on_record(txn, rid, tab_fd, LockMode::SHARED);
}

/**
 * @description: 申请行级排他锁
 * @return {bool} 加锁是否成功
 * @param {Transaction*} txn 要申请锁的事务对象指针
 * @param {Rid&} rid 加锁的目标记录ID
 * @param {int} tab_fd 记录所在的表的fd
 */
bool LockManager::lock_exclusive_on_record(Transaction* txn, const Rid& rid, int tab_fd) {
    return lock_on_record(txn, rid, tab_fd, LockMode::EXCLUSIVE);
}

/**
 * @description: 申请表级读锁
 * @return {bool} 返回加锁是否成功
 * @param {Transaction*} txn 要申请锁的事务对象指针
 * @param {int} tab_fd 目标表的fd
 */
bool LockManager::lock_shared_on_table(Transaction* txn, int tab_fd) {
    return lock_on_table(txn, tab_fd, LockMode::SHARED);
}

/**
 * @description: 申请表级写锁
 * @return {bool} 返回加锁是否成功
 * @param {Transaction*} txn 要申请锁的事务对象指针
 * @param {int} tab_fd 目标表的fd
 */
bool LockManager::lock_exclusive_on_table(Transaction* txn, int tab_fd) {
    return lock_on_table(txn, tab_fd, LockMode::EXCLUSIVE);
}

/**
 * @description: 申请表级意向读锁
 * @return {bool} 返回加锁是否成功
 * @param {Transaction*} txn 要申请锁的事务对象指针
 * @param {int} tab_fd 目标表的fd
 */
bool LockManager::lock_IS_on_table(Transaction* txn, int tab_fd) {
    return lock_on_table(txn, tab_fd, LockMode::INTENTION_SHARED);
}

/**
 * @description: 申请表级意向写锁
 * @return {bool} 返回加锁是否成功
 * @param {Transaction*} txn 要申请锁的事务对象指针
 * @param {int} tab_fd 目标表的fd
 */
bool LockManager::lock_IX_on_table(Transaction* txn, int tab_fd) {
    return lock_on_table(txn, tab_fd, LockMode::INTENTION_EXCLUSIVE);
}

/**
 * @description: 释放锁
 * @return {bool} 返回解锁是否成功
 * @param {Transaction*} txn 要释放锁的事务对象指针
 * @param {LockDataId} lock_data_id 要释放的锁ID
 */
bool LockManager::unlock(Transaction* txn, LockDataId lock_data_id) {
    if (lock_data_id.type_ == LockDataType::RECORD) {
        return unlock_record(txn, lock_data_id);
    } else {
        return unlock_table(txn, lock_data_id);
    }
}

bool LockManager::lock_on_record(Transaction * txn, const Rid & rid, int tab_fd, LockMode lock_mode) {
    // 检查事务状态
    assert(txn->get_isolation_level() == IsolationLevel::SERIALIZABLE);
    if (txn->get_state() == TransactionState::DEFAULT) {
        txn->set_state(TransactionState::GROWING);
    }
    if (txn->get_state() != TransactionState::GROWING) {
        return false;
    }

    LockDataId lock_data_id(tab_fd, rid, LockDataType::RECORD);
    bool upgrade = false;
    if (lock_mode == LockMode::SHARED) {
        if (txn->is_s_rec_locked(lock_data_id) || txn->is_x_rec_locked(lock_data_id)) {
            return true;
        }
    } else {
        if (txn->is_x_rec_locked(lock_data_id)) {
            return true;
        }
        if (txn->is_s_rec_locked(lock_data_id)) {
            upgrade = true;
        }
    }

    latch_.lock();
    auto lock_table_it = lock_table_.find(lock_data_id);

    auto lock_request = std::make_shared<LockRequest>(txn->get_transaction_id(), lock_mode);

    if (lock_table_it != lock_table_.end()) {
        auto &lrque = lock_table_it->second;
        

        auto &request_queue = lrque->request_queue_;
        decltype(request_queue.begin()) request_it;

        if (upgrade) {
            if (lrque->upgrading_) {
                txn->set_state(TransactionState::ABORTED);
                throw TransactionAbortException(txn->get_transaction_id(), AbortReason::UPGRADE_CONFLICT);
            }

            request_it = std::find_if(request_queue.begin(), request_queue.end(), 
                                    [txn](const auto &req) { return req->txn_id_ == txn->get_transaction_id();});
            lrque->upgrading_ = true;

            auto it = std::find_if(request_it, request_queue.end(), [](const auto &req) { return !req->granted_;});
            request_queue.erase(request_it);
            request_it = request_queue.insert(it, lock_request);
            lrque->update_group_lock_mode();

            txn->get_s_rec_lock_set()->erase(lock_data_id);
        } else {
            request_it = request_queue.insert(request_queue.end(), lock_request);
        }

        // 检查相容性看是否能够授予锁
        if (!lrque->check_compatible(lock_mode)) {
            txn->set_state(TransactionState::ABORTED);
            request_queue.erase(request_it);
            if (upgrade) lrque->upgrading_ = false;
            latch_.unlock();
            throw TransactionAbortException(txn->get_transaction_id(), AbortReason::DEADLOCK_PREVENTION);
        }

        lock_request->granted_ = true;
        lrque->update_group_lock_mode();
        if (upgrade) lrque->upgrading_ = false;
        latch_.unlock();
    } else {
        lock_table_.insert(std::make_pair(lock_data_id, std::make_shared<LockRequestQueue>()));
        lock_table_[lock_data_id]->request_queue_.push_back(lock_request);
        lock_request->granted_ = true;
        lock_table_[lock_data_id]->update_group_lock_mode();
        latch_.unlock();
    }

    if (lock_mode == LockMode::SHARED) {
        txn->get_s_rec_lock_set()->insert(lock_data_id);
    } else {
        txn->get_x_rec_lock_set()->insert(lock_data_id);
    }
    return true;
}

bool LockManager::lock_on_table(Transaction * txn, int tab_fd, LockMode lock_mode) {
    assert(txn->get_isolation_level() == IsolationLevel::SERIALIZABLE);
    // 检查事务的状态
    if (txn->get_state() == TransactionState::DEFAULT) {
        txn->set_state(TransactionState::GROWING);
    }
    if (txn->get_state() != TransactionState::GROWING) {
        return false;
    }

    LockDataId lock_data_id(tab_fd, LockDataType::TABLE);
    auto old_lock_mode = is_table_locked(txn, lock_data_id, ANYLOCK);
    bool upgrade = false;
    if (old_lock_mode.has_value()) {
        // 当持有相同的锁，或者持有更高级别（排他性更强）的锁时，直接返回
        if (old_lock_mode.value() == lock_mode || is_stronger(old_lock_mode.value(), lock_mode)) {
            return true;
        }
        switch (old_lock_mode.value()) {
            case LockMode::INTENTION_SHARED:
                break;
            case LockMode::SHARED:
            case LockMode::INTENTION_EXCLUSIVE:
                if (lock_mode != LockMode::EXCLUSIVE && lock_mode != LockMode::S_IX) {
                    txn->set_state(TransactionState::ABORTED);
                    throw TransactionAbortException(txn->get_transaction_id(), AbortReason::DEADLOCK_PREVENTION);
                }
                break;
            case LockMode::S_IX:
                if (lock_mode != LockMode::EXCLUSIVE) {
                    txn->set_state(TransactionState::ABORTED);
                    throw TransactionAbortException(txn->get_transaction_id(), AbortReason::DEADLOCK_PREVENTION);
                }
                break;
            case LockMode::EXCLUSIVE:
                txn->set_state(TransactionState::ABORTED);
                throw TransactionAbortException(txn->get_transaction_id(), AbortReason::DEADLOCK_PREVENTION);
                break;
        }
        upgrade = true;
    }

    latch_.lock();
    auto lock_table_it = lock_table_.find(lock_data_id);
    auto lock_request = std::make_shared<LockRequest>(txn->get_transaction_id(), lock_mode);

    if (lock_table_it != lock_table_.end()) {
        auto &lrque = lock_table_it->second;
        auto &request_queue = lrque->request_queue_;
        decltype(request_queue.begin()) request_it;

        if (upgrade) {
            if (lrque->upgrading_) {
                txn->set_state(TransactionState::ABORTED);
                throw TransactionAbortException(txn->get_transaction_id(), AbortReason::UPGRADE_CONFLICT);
            }
            request_it = std::find_if(request_queue.begin(), request_queue.end(),
                                      [&txn](const auto &lr) { return lr->txn_id_ == txn->get_transaction_id(); });
            
            lrque->upgrading_ = true;

            auto it = std::find_if(request_it, request_queue.end(), [](const auto &lr) { return !lr->granted_; });
            request_queue.erase(request_it);
            request_queue.insert(it, lock_request);
            lrque->update_group_lock_mode();

            switch (old_lock_mode.value()) {
                case LockMode::SHARED:
                    txn->get_s_tab_lock_set()->erase(lock_data_id);
                    break;
                case LockMode::EXCLUSIVE:
                    txn->get_x_tab_lock_set()->erase(lock_data_id);
                    break;
                case LockMode::INTENTION_SHARED:
                    txn->get_is_tab_lock_set()->erase(lock_data_id);
                    break;
                case LockMode::INTENTION_EXCLUSIVE:
                    txn->get_ix_tab_lock_set()->erase(lock_data_id);
                    break;
                case LockMode::S_IX:
                    txn->get_six_tab_lock_set()->erase(lock_data_id);
                    break;
            }
        } else {
            request_it = request_queue.insert(request_queue.end(), lock_request);
        }

        if (!lrque->check_compatible(lock_mode)) {
            txn->set_state(TransactionState::ABORTED);
            request_queue.erase(request_it);
            if (upgrade) lrque->upgrading_ = false;
            latch_.unlock();
            throw TransactionAbortException(txn->get_transaction_id(), AbortReason::DEADLOCK_PREVENTION);
        }

        lock_request->granted_ = true;
        lrque->update_group_lock_mode();
        if (upgrade) lrque->upgrading_ = false;
        latch_.unlock();
    } else {
        lock_table_.insert(std::make_pair(lock_data_id, std::make_shared<LockRequestQueue>()));
        lock_table_[lock_data_id]->request_queue_.push_back(lock_request);
        lock_request->granted_ = true;
        lock_table_[lock_data_id]->update_group_lock_mode();
        latch_.unlock();
    }

    switch (lock_mode) {
        case LockMode::SHARED:
            txn->get_s_tab_lock_set()->insert(lock_data_id);
            break;
        case LockMode::EXCLUSIVE:
            txn->get_x_tab_lock_set()->insert(lock_data_id);
            break;
        case LockMode::INTENTION_SHARED:
            txn->get_is_tab_lock_set()->insert(lock_data_id);
            break;
        case LockMode::INTENTION_EXCLUSIVE:
            txn->get_ix_tab_lock_set()->insert(lock_data_id);
            break;
        case LockMode::S_IX:
            txn->get_six_tab_lock_set()->insert(lock_data_id);
            break;
    }
    return true;
}

void LockManager::unlock_all(Transaction* txn) {
    // 释放表锁
    while(! txn->get_s_tab_lock_set()->empty()) {
        auto lock_data_id = *txn->get_s_tab_lock_set()->begin();
        unlock_table(txn, lock_data_id);
    }
    while(! txn->get_x_tab_lock_set()->empty()) {
        auto lock_data_id = *txn->get_x_tab_lock_set()->begin();
        unlock_table(txn, lock_data_id);
    }
    while(! txn->get_is_tab_lock_set()->empty()) {
        auto lock_data_id = *txn->get_is_tab_lock_set()->begin();
        unlock_table(txn, lock_data_id);
    }
    while(! txn->get_ix_tab_lock_set()->empty()) {
        auto lock_data_id = *txn->get_ix_tab_lock_set()->begin();
        unlock_table(txn, lock_data_id);
    }
    while(! txn->get_six_tab_lock_set()->empty()) {
        auto lock_data_id = *txn->get_six_tab_lock_set()->begin();
        unlock_table(txn, lock_data_id);
    }
    // 释放行锁
    while(! txn->get_s_rec_lock_set()->empty()) {
        auto lock_data_id = *txn->get_s_rec_lock_set()->begin();
        unlock_record(txn, lock_data_id);
    }
    while(! txn->get_x_rec_lock_set()->empty()) {
        auto lock_data_id = *txn->get_x_rec_lock_set()->begin();
        unlock_record(txn, lock_data_id);
    }
}

bool LockManager::unlock_record(Transaction* txn, LockDataId lock_data_id) {
    // 检查事务状态
    if (txn->get_state() < TransactionState::SHRINKING) {
        txn->set_state(TransactionState::SHRINKING);
    }
    // 不再判断，因为有可能是事务回滚导致的Unlock
    // if (txn->get_state() != TransactionState::SHRINKING) {
    //     return false;
    // }
    latch_.lock();
    auto lock_table_it = lock_table_.find(lock_data_id);
    assert(lock_table_it != lock_table_.end());
    LockMode lock_mode;
    auto &lrque = lock_table_it->second;

    auto &request_queue = lrque->request_queue_;
    auto it = std::find_if(request_queue.begin(), request_queue.end(),
                        [txn](const auto &request) { return request->txn_id_ == txn->get_transaction_id(); });
    assert(it != request_queue.end() && it->get()->granted_);

    lock_mode = it->get()->lock_mode_;
    request_queue.erase(it);
    lrque->update_group_lock_mode();
    latch_.unlock();
    
    switch(lock_mode) {
        case LockMode::SHARED:
            txn->get_s_rec_lock_set()->erase(lock_data_id);
            break;
        case LockMode::EXCLUSIVE:
            txn->get_x_rec_lock_set()->erase(lock_data_id);
            break;
        default:
            assert(false);
    }
    return true;
}

bool LockManager::unlock_table(Transaction* txn, LockDataId lock_data_id) {
    // 检查事务状态
    if (txn->get_state() < TransactionState::SHRINKING) {
        txn->set_state(TransactionState::SHRINKING);
    }
    // 不再判断，因为有可能是事务回滚导致的Unlock
    // if (txn->get_state() != TransactionState::SHRINKING) {
    //     return false;
    // }
    latch_.lock();
    auto lock_table_it = lock_table_.find(lock_data_id);
    assert(lock_table_it != lock_table_.end());
    LockMode lock_mode;
    auto &lrque = lock_table_it->second;

    auto &request_queue = lrque->request_queue_;
    auto it = std::find_if(request_queue.begin(), request_queue.end(),
                        [txn](const auto &request) { return request->txn_id_ == txn->get_transaction_id(); });
    assert(it != request_queue.end() && it->get()->granted_ == true);

    lock_mode = it->get()->lock_mode_;
    request_queue.erase(it);
    lrque->update_group_lock_mode();
    latch_.unlock();

    switch (lock_mode) {
        case LockMode::SHARED:
            txn->get_s_tab_lock_set()->erase(lock_data_id);
            break;
        case LockMode::EXCLUSIVE:
            txn->get_x_tab_lock_set()->erase(lock_data_id);
            break;
        case LockMode::INTENTION_SHARED:
            txn->get_is_tab_lock_set()->erase(lock_data_id);
            break;
        case LockMode::INTENTION_EXCLUSIVE:
            txn->get_ix_tab_lock_set()->erase(lock_data_id);
            break;
        case LockMode::S_IX:
            txn->get_six_tab_lock_set()->erase(lock_data_id);
            break;
    }
    return true;
}

std::optional<LockManager::LockMode> LockManager::is_table_locked(Transaction* txn, const LockDataId& lock_data_id, const std::vector<LockMode>& lock_modes) {
    std::optional<LockMode> mode = std::nullopt;
    for (const auto &lock_mode: lock_modes) {
        switch(lock_mode) {
            case LockMode::SHARED:
                if (txn->is_s_tab_locked(lock_data_id)) {
                    mode = std::make_optional<LockMode>(LockMode::SHARED);
                }
                break;
            case LockMode::EXCLUSIVE:
                if (txn->is_x_tab_locked(lock_data_id)) {
                    mode = std::make_optional<LockMode>(LockMode::EXCLUSIVE);
                }
                break;
            case LockMode::INTENTION_SHARED:
                if (txn->is_is_tab_locked(lock_data_id)) {
                    mode = std::make_optional<LockMode>(LockMode::INTENTION_SHARED);
                }
                break;
            case LockMode::INTENTION_EXCLUSIVE:
                if (txn->is_ix_tab_locked(lock_data_id)) {
                    mode = std::make_optional<LockMode>(LockMode::INTENTION_EXCLUSIVE);
                }
                break;
            case LockMode::S_IX:
                if (txn->is_six_tab_locked(lock_data_id)) {
                    mode = std::make_optional<LockMode>(LockMode::S_IX);
                }
                break;
        }
        if (mode.has_value()) {
            break;
        }
    }
    return mode;
}

// 判断 old_lock_mode 的排他性是否比 new_lock_mode 更强
bool LockManager::is_stronger(LockMode old_lock_mode, LockMode new_lock_mode) {
    assert(old_lock_mode != new_lock_mode);
    if (old_lock_mode == LockMode::EXCLUSIVE) {
        return true;
    } else if (old_lock_mode == LockMode::S_IX && new_lock_mode != LockMode::EXCLUSIVE) {
        return true;
    } else if (old_lock_mode == LockMode::INTENTION_EXCLUSIVE && new_lock_mode == LockMode::INTENTION_SHARED) {
        return true;
    } else if (old_lock_mode == LockMode::SHARED && new_lock_mode == LockMode::INTENTION_SHARED) {
        return true;
    } else {
        return false;
    }
}

std::string LockManager::lockmode2str(LockMode lock_mode) {
    switch (lock_mode) {
        case LockMode::SHARED: {
            return "SHARED";
        } break;

        case LockMode::EXCLUSIVE: {
            return "EXCLUSIVE";
        } break;

        case LockMode::INTENTION_SHARED: {
            return "INTENTION_SHARED";
        } break;

        case LockMode::INTENTION_EXCLUSIVE: {
            return "INTENTION_EXCLUSIVE";
        } break;

        case LockMode::S_IX: {
            return "S_IX";
        } break;

        default: {
            assert(false);
        } break;
    }
}

 /* 根据组模式，检查是否可以授予锁
        Compatibility Matrix
         IS   IX   S   SIX   X
    IS    √    √    √    √    ×
    IX    √    √    ×    ×    ×
    S     √    ×    √    ×    ×
    SIX   √    ×    ×    ×    ×
    X     ×    ×    ×    ×    ×
*/
bool LockManager::LockRequestQueue::check_compatible(LockMode lock_mode) {
    switch (lock_mode) {
        case LockMode::INTENTION_SHARED: {
            return group_lock_mode_ == GroupLockMode::NON_LOCK || group_lock_mode_ == GroupLockMode::IS ||
                    group_lock_mode_ == GroupLockMode::IX || group_lock_mode_ == GroupLockMode::S || 
                    group_lock_mode_ == GroupLockMode::SIX;
        } break;

        case LockMode::INTENTION_EXCLUSIVE: {
            return group_lock_mode_ == GroupLockMode::NON_LOCK || group_lock_mode_ == GroupLockMode::IS ||
                    group_lock_mode_ == GroupLockMode::IX;
        } break;

        case LockMode::SHARED: {
            return group_lock_mode_ == GroupLockMode::NON_LOCK || group_lock_mode_ == GroupLockMode::IS ||
                    group_lock_mode_ == GroupLockMode::S;
        } break;

        case LockMode::S_IX: {
            return group_lock_mode_ == GroupLockMode::NON_LOCK || group_lock_mode_ == GroupLockMode::IS;
        } break;

        case LockMode::EXCLUSIVE: {
            return group_lock_mode_ == GroupLockMode::NON_LOCK;
        } break;

        default: {
            assert(false);
        } break;
    }
}

void LockManager::LockRequestQueue::update_group_lock_mode() {
    // 遍历，查看已经授予的锁的类型
    //修改后的组模式
    GroupLockMode mode = GroupLockMode::NON_LOCK;
    for (auto &req_iter : request_queue_) {
        if (!req_iter->granted_) continue;
        if (req_iter->lock_mode_ == LockMode::EXCLUSIVE){
            mode = GroupLockMode::X;
            break;
        } else if (req_iter->lock_mode_ == LockMode::SHARED && mode != GroupLockMode::SIX) {
            if (mode == GroupLockMode::IX) {
                mode = GroupLockMode::SIX;
            } else {
                mode = GroupLockMode::S;
            }
        } else if (req_iter->lock_mode_ == LockMode::INTENTION_EXCLUSIVE && mode != GroupLockMode::SIX) {
            if (mode == GroupLockMode::S){
                mode = GroupLockMode::SIX;
            } else {
                mode = GroupLockMode::IX;
            }
        } else if (req_iter->lock_mode_ == LockMode::INTENTION_SHARED &&
                (mode == GroupLockMode::NON_LOCK || mode == GroupLockMode::IS)) {
            mode = GroupLockMode::IS;
        }
    }
    group_lock_mode_ = mode;
}